% author: ziyan zhu (zzhu1@g.harvard.edu) 
% load data from Julia calculations, plot the local misfit energy
% distribution and local twist angle distributions
clear all 
close all
f_size = 22;
set(groot, 'DefaultTextInterpreter', 'Latex')
set(groot, 'DefaultLegendInterpreter', 'Latex')
set(groot, 'DefaultAxesTickLabelInterpreter', 'Latex')
set(0,'DefaultAxesFontSize',f_size)

system = 'triG';

scale = 15; 
N = 54; % discretization. change accordingly 

q12 = deg2rad(1.50); % twist angles
q23 = deg2rad(1.69);

% load data from Julia calculations
masterdir = '/data_amp/';

curl = importdata(['.' masterdir 'triG_q12_' num2str(rad2deg(q12)) 'deg_q23_' num2str(rad2deg(q23))...
    'deg_N_' num2str(N) '_scale_' num2str(scale) '_curl.txt']);
energy_data = importdata(['.' masterdir 'triG_q12_' num2str(rad2deg(q12),'%.1f') 'deg_q23_' num2str(rad2deg(q23))...
    'deg_N_' num2str(N) '_scale_' num2str(scale) '_disp_energy.txt']);
xarr = importdata(['.' masterdir 'triG_q12_' num2str(rad2deg(q12),'%.1f') 'deg_q23_' num2str(rad2deg(q23)) ...
    'deg_N_' num2str(N) '_scale_' num2str(scale) '_xarr.txt']);
yarr = importdata(['.' masterdir 'triG_q12_' num2str(rad2deg(q12),'%.1f') 'deg_q23_' num2str(rad2deg(q23)) ...
    'deg_N_' num2str(N) '_scale_' num2str(scale) '_yarr.txt']);


%%
x_energy = energy_data.data(:, 1);
y_energy = energy_data.data(:, 2);
energy = energy_data.data(:,3:4);
[xmesh, ymesh] = meshgrid(xarr(:), yarr(:));

curl = -reshape(curl, [length(xarr), length(yarr), 3]);
curl = permute(curl, [2, 1, 3]);

%%

a0 = 1.42*sqrt(3);

rot12 = [cos(q12), sin(q12); -sin(q12), cos(q12)];
rot23 = [cos(q23), sin(q23); -sin(q23), cos(q23)];

asc_12 = a0/(2*sin(q12/2));
Asc_12 = [0, 0;
        -0.5*asc_12, sqrt(3)/2*asc_12;
        0, sqrt(3)*asc_12;
        0.5*asc_12, sqrt(3)/2*asc_12];
asc_23 = a0/(2*sin(q23/2));
Asc_23 = [0, 0;
        -0.5*asc_23, sqrt(3)/2*asc_23;
        0, sqrt(3)*asc_23;
        0.5*asc_23, sqrt(3)/2*asc_23];
    
    
for i = 1:size(Asc_12,1)
    Asc_12(i, :) = (rot12 *inv(rot23) * Asc_12(i, :)')';
    Asc_23(i, :) = (rot23 *inv(rot12)* Asc_23(i, :)')';
end 

Asc_x = [Asc_12(:, 1); Asc_23(:, 1)];
Asc_y = [Asc_12(:, 2); Asc_23(:, 2)];


s_grid = 5; 
vec = [sqrt(3)/2    sqrt(3)/2;
       -1/2          1/2];
   

A1 = rot12 * a0 * vec; 
A2 = a0 * vec;
A3 = rot23 * a0 * vec; 


%% energy profile 
n = 500;

X = linspace(min(xarr), max(xarr), 4*n);
Y = linspace(min(yarr), max(yarr), 5*n);

xtmp = X' .* ones(1,5*n);
ytmp = ones(4*n,1) .* reshape(Y, [1,5*n]);

misfit12 = reshape(energy(:,1), size(xtmp))*1e3;
misfit23 = reshape(energy(:,2), size(ytmp))*1e3;

% downsampling 
clear xy 
int = 1;
misfit12_down = misfit12(1:int:end, 1:int:end); 
misfit23_down = misfit23(1:int:end, 1:int:end);
xarr_down = xtmp(1:int:end, 1:int:end);
yarr_down = ytmp(1:int:end, 1:int:end);

misfit_down = misfit12_down + misfit23_down;


%% plot the relaxed misfit energy as a function of space 

xl2 = [-650 650];
yl2 = [-200 1500];
    
figure 
set(gcf, 'Position', [45 270 1313 493]);
tiledlayout(1,3)
for i = 1:3
    nexttile; 
    hold all
    if i == 1
        surf( xarr_down, yarr_down, misfit12_down, 'EdgeColor', 'none')
        ylabel("$y$ (\AA)")
        title("$E_{12}$")
        caxis([0 max(misfit12(:))])
    elseif i == 2 
        surf( xarr_down, yarr_down, misfit23_down, 'EdgeColor', 'none')
        title("$E_{23}$")
        caxis([0 max(misfit23(:))])
    elseif i == 3 
        surf( xarr_down, yarr_down, misfit12_down+misfit23_down, 'EdgeColor', 'none')
        title("$E_\mathrm{tot}$")
        caxis([0 max(misfit12(:)+misfit23(:))])
    end 
    view(2)
    xlabel("$x$ (\AA)");
    c1 = colorbar('southoutside');
    c1.TickLabelInterpreter= 'latex';
    cmap=colormap(brewermap([],'RdBu'));
    cmap=flip(cmap,1);
    colormap(cmap)
    xlim(xl2)
    ylim(yl2)
    axis equal
    
end 



%% calculate & plot local twist angle 

theta12 = rad2deg(abs(q12 + asin(curl(:,:,2)/2) + asin(curl(:,:,1)/2)));
theta23 = rad2deg(abs(q23 + asin(curl(:,:,3)/2) + asin(curl(:,:,2)/2)));
N=200;
[rgb,num,typ] = brewermap(N,'RdBu');


figure
set(gcf, 'Position', [120 235 1117 515])
tiledlayout(1,3)
for i = 1:3
    nexttile; 

    if i == 1
        surf(xmesh, ymesh, theta12, 'EdgeColor', 'none')
        ylabel("$y$ (\AA)")
        ct_label = "$\theta_{TM, \mathrm{local}} (^\circ)$";

    elseif i == 2 
        surf(xmesh, ymesh, theta23, 'EdgeColor', 'none')
        ct_label = "$\theta_{MB, \mathrm{local}} (^\circ)$";

    elseif i == 3 
        surf(xmesh, ymesh, theta12-theta23, 'EdgeColor', 'none')
        ct_label = "$\theta_{TM, \mathrm{local}}-\theta_{MB, \mathrm{local}} (^\circ)$";
    end 
    axis equal 
    box on

    if i == 2 
        xlabel("$x$ (\AA)")
    else 
        xlabel("$\quad$")
    end 


    c1 = colorbar('southoutside');
    c1.TickLabelInterpreter= 'latex';

    view(2)
    xlim(xl2*1.2)
    ylim(yl2*1.2)

    xl_here = xlim;
    yl_here = ylim;
%         xlabel(c1,ct_label);
    if i < 3
        cx = (max(xl_here)-min(xl_here))*0.29+min(xl_here);
        caxis([0 max([theta12(:);theta23(:)])])

    else 
        cx = (max(xl_here)-min(xl_here))*0.15+min(xl_here); 
        caxis([-1, 1].*max(abs(theta12(:)-theta23(:))))
    end 
    cy = min(yl_here)*3;
    text(cx, cy, ct_label, 'FontSize', f_size);

    colormap(rgb)
%         caxis([1.4, 1.9])
end 
